import mindspore as ms
import mindspore.dataset as ds
from .data_utils import get_rank_info, get_num_parallel_workers


def create_dataset(dataset_path, train, batch_size=32, train_image_size=224, distribute=True):
    device_num, rank_id = get_rank_info(distribute)
    ds.config.set_prefetch_size(64)
    if train:
        data_set = ds.Cifar10Dataset(dataset_path, usage='train', num_parallel_workers=get_num_parallel_workers(12), shuffle=True, num_shards=device_num, shard_id=rank_id)
    else:
        data_set = ds.Cifar10Dataset(dataset_path, usage='test', num_parallel_workers=get_num_parallel_workers(12), shuffle=False, num_shards=device_num, shard_id=rank_id)

    # define map operations
    trans = []
    if train:
        trans += [
            ds.vision.c_transforms.RandomCrop((32, 32), (4, 4, 4, 4)),
            ds.vision.c_transforms.RandomHorizontalFlip(prob=0.5)
        ]

    trans += [
        ds.vision.c_transforms.Resize((train_image_size, train_image_size)),
        ds.vision.c_transforms.Rescale(1.0 / 255.0, 0.0),
        ds.vision.c_transforms.Normalize([0.4914, 0.4822, 0.4465], [0.2023, 0.1994, 0.2010]),
        ds.vision.c_transforms.HWC2CHW()
    ]

    type_cast_op = ds.transforms.c_transforms.TypeCast(ms.int32)

    data_set = data_set.map(operations=type_cast_op, input_columns="label", num_parallel_workers=get_num_parallel_workers(8))
    # only enable cache for eval
    data_set = data_set.map(operations=trans, input_columns="image", num_parallel_workers=get_num_parallel_workers(8))

    # apply batch operations
    data_set = data_set.batch(batch_size, drop_remainder=True)

    return data_set
